import pymysql
import pandas as pd


class MysqlUtils:
    def __init__(self):
        try:
            self.conn = pymysql.connect(
                host='127.0.0.1',
                user='root',
                passwd='sjk1234',
                db='msy',
                port=3306,
                charset='utf8'
            )
        except pymysql.Error as e:
            print(f"数据库连接失败: {e}")

    def __del__(self):
        if hasattr(self, 'conn') and self.conn:
            self.conn.close()


class Classification:
    def get_fina_indicator(self, conn):
        """
        获取财务数据
        """
        try:
            with conn.cursor(cursor=pymysql.cursors.DictCursor) as cursor:
                sql = """
                SELECT ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps, gross_margin, fcff, fcf
                FROM your_table_name  -- 请替换为实际表名
                """
                cursor.execute(sql)
                ret = cursor.fetchall()
            df = pd.DataFrame(ret)
            # 数据清洗
            df = df.dropna(subset=['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcf'])
            # 重建索引
            df = df.reset_index(drop=True)
            print(df.head().to_csv(sep='\t', na_rep='nan'))
            return df
        except pymysql.Error as e:
            print(f"获取财务数据时出错: {e}")
            return pd.DataFrame()

    def get_daily(self, conn, df):
        try:
            new_list = []
            with conn.cursor(cursor=pymysql.cursors.DictCursor) as cursor:
                for _, row in df.iterrows():
                    ann_date_str = row['ann_date'].strftime('%Y%m%d')
                    sql = f"SELECT trade_date, closes FROM date_1 WHERE ts_code = '{row['ts_code']}'"
                    cursor.execute(sql)
                    ret = cursor.fetchall()
                    for r in ret:
                        data_dict = {
                            'trade_date': r['trade_date'],
                            'closes': r['closes'],
                            'the_slope': '',
                            'eps': row['eps'],
                            'total_revenue_ps': row['total_revenue_ps'],
                            'undist_profit_ps': row['undist_profit_ps'],
                            'gross_margin': row['gross_margin'],
                            'fcff': row['fcff'],
                            'fcfe': row.get('fcfe', None),
                            'tangible_asset': row.get('tangible_asset', None),
                            'bps': row.get('bps', None),
                            'grossprofit_margin': row.get('grossprofit_margin', None),
                            'npta': row.get('npta', None)
                        }
                        new_list.append(data_dict)
            df2 = pd.DataFrame(new_list)
            df2.to_csv('daily.csv', index=False)
        except pymysql.Error as e:
            print(f"获取每日数据时出错: {e}")


if __name__ == '__main__':
    mu = MysqlUtils()
    ci = Classification()
    if hasattr(mu, 'conn') and mu.conn:
        df = ci.get_fina_indicator(mu.conn)
        if not df.empty:
            ci.get_daily(mu.conn, df)